#include "soft_min_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
#include <vector>
#include <cstdint>
#include <cmath>
#include <random>
#include <iostream>
#include <algorithm>

namespace optiling {
    static inline uint32_t AlignUp(uint32_t a, uint32_t b) 
    {
        if (b == 0)
            return a;
        return (a + b - 1) / b * b;
    }
    
    const uint32_t BLOCK_SIZE = 32;         // 基础对齐块大小（字节）
    const uint32_t CORE_NUM = 8;            // AI Core 数量
    const uint32_t MIN_TILE_SIZE = 16;      // 最小切分大小
    const uint32_t BUFFER_NUM = 2;          // 双缓冲数量
    const uint32_t VEC_LEN = 8;             // 向量化长度

    static ge::graphStatus TilingFunc(gert::TilingContext* context) {
        SoftMinCustomTilingData tiling;
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        
        // 设置多核并行
        context->SetBlockDim(CORE_NUM);

        // 获取输入张量形状和属性
        const auto InputShape = context->GetInputTensor(0)->GetOriginShape();
        const auto inputDataType = context->GetInputTensor(0)->GetDataType();
        
        uint32_t typeSize = sizeof(float);
        
        // 获取dim参数
        int32_t dim = 1; // 默认dim=1 (按行计算)
        auto attrPtr = context->GetAttrs();
        if (attrPtr && attrPtr->GetInt(0)) 
        {
            dim = static_cast<int32_t>(*(attrPtr->GetInt(0)));
        }
        
        // 设置tilingkey
        switch (dim) {
            case 1:
                context->SetTilingKey(1);
                break;
            case 2:
                context->SetTilingKey(2);
                break;
            default:
                return ge::GRAPH_FAILED;
        }

        // 修正：检查张量维度数
        auto dimNum = InputShape.GetDimNum();
        if (dimNum < 2) {
            return ge::GRAPH_FAILED;
        }

        // 计算各维度大小
        int32_t batchSize = 1;
        for (int i = 0; i < dimNum - 2; i++) {
            batchSize *= InputShape.GetDim(i);
        }
        int32_t rows = InputShape.GetDim(dimNum - 2);  // 序列长度(行数)
        int32_t cols = InputShape.GetDim(dimNum - 1);  // 特征维度(列数) 

        int32_t computeSize = (dim == 1) ? cols : rows; // 每次softmin计算的向量长度
        int32_t vectorNum = (dim == 1) ? rows : cols;   // 需要计算的向量数量

        // 获取UB大小
        uint64_t ubSize = 256 * 1024; // 默认256KB
        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);

        // 核间划分策略
        int32_t total_vectors = batchSize * vectorNum;
        
        // 修正：处理边界情况
        if (total_vectors == 0) {
            return ge::GRAPH_FAILED;
        }
        
        int32_t base_vectors_per_core = total_vectors / CORE_NUM;
        int32_t big_core_num = total_vectors % CORE_NUM; // 大核个数
        int32_t big_core_vectors = base_vectors_per_core + 1; // 大核处理的向量数
        int32_t small_core_num = CORE_NUM - big_core_num; // 小核个数
        int32_t small_core_vectors = base_vectors_per_core; // 小核处理的向量数

        // 分配数组存储每个核的划分信息
        int32_t* core_vector_start = new int32_t[CORE_NUM];
        int32_t* core_vector_end = new int32_t[CORE_NUM];
        int32_t* core_vector_count = new int32_t[CORE_NUM];

        auto cleanup = [&]() {
            delete[] core_vector_start;
            delete[] core_vector_end;
            delete[] core_vector_count;
        };

        // 计算每个核的向量分配
        int32_t current_vector = 0;
        for (int32_t core_id = 0; core_id < CORE_NUM; core_id++) {
            if (core_id < big_core_num) {
                core_vector_count[core_id] = big_core_vectors;
            } else {
                core_vector_count[core_id] = small_core_vectors;
            }
            core_vector_start[core_id] = current_vector;
            
            if (core_vector_count[core_id] > 0) {
                core_vector_end[core_id] = current_vector + core_vector_count[core_id] - 1;
                current_vector += core_vector_count[core_id];
            } 
            else {
                core_vector_end[core_id] = current_vector - 1; // 处理空核的情况
            }
        }

        // 核内划分策略：简化为一行一行搬运
        int32_t* core_loop_times = new int32_t[CORE_NUM];

        for (int core_id = 0; core_id < CORE_NUM; core_id++) {
            if (core_vector_count[core_id] > 0) {
                core_loop_times[core_id] = core_vector_count[core_id]; // 每个核循环处理的次数
            } 
            else {
                core_loop_times[core_id] = 0;
            }
        }


        auto cleanup_all = [&]() {
            cleanup();
            delete[] core_loop_times;
        };

        // 计算对齐参数
        int32_t aligned_compute_size = ((computeSize + 32 / (int32_t)typeSize - 1) / (32 / (int32_t)typeSize)) * (32 / (int32_t)typeSize);

        // 设置tiling数据
        tiling.set_batchSize(static_cast<int32_t>(batchSize));
        tiling.set_rows(static_cast<int32_t>(rows));
        tiling.set_cols(static_cast<int32_t>(cols));
        tiling.set_computeSize(static_cast<int32_t>(computeSize));
        tiling.set_vectorNum(static_cast<int32_t>(vectorNum));
        tiling.set_dim(dim);

        // 核间划分信息
        tiling.set_total_vectors(total_vectors);
        tiling.set_base_vectors_per_core(base_vectors_per_core);
        tiling.set_big_core_num(big_core_num);
        tiling.set_core_vector_start(core_vector_start);
        tiling.set_core_vector_end(core_vector_end);
        tiling.set_core_vector_count(core_vector_count);

        // 核内划分信息
        tiling.set_core_loop_times(core_loop_times);
        tiling.set_aligned_compute_size(aligned_compute_size);

        // 设置工作空间大小
        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        if (!currentWorkspace) { 
            cleanup_all();
            return ge::GRAPH_FAILED;
        }
        currentWorkspace[0] = 0;

        // 保存tiling数据
        auto tilingData = context->GetRawTilingData();
        if (!tilingData) {
            cleanup_all();
            return ge::GRAPH_FAILED;
        }
        
        tiling.SaveToBuffer(tilingData->GetData(), tilingData->GetCapacity());
        tilingData->SetDataSize(tiling.GetDataSize());
        
        cleanup_all();
        return ge::GRAPH_SUCCESS;
    }
}

namespace ge {
static graphStatus InferShape(gert::InferShapeContext *context)
{
    if (!context) {
        return ge::GRAPH_FAILED;
    }

    const gert::Shape *InputShape = context->GetInputShape(0);
    gert::Shape *OutputShape = context->GetOutputShape(0);
    
    if (!InputShape || !OutputShape) {
        return ge::GRAPH_FAILED;
    }

    // 直接复制输入形状到输出
    *OutputShape = *InputShape;

    return ge::GRAPH_SUCCESS;
}

static graphStatus InferDataType(gert::InferDataTypeContext *context)
{
    if (!context) {
        return ge::GRAPH_FAILED;
    }

    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}

} // namespace ge

namespace ops {
class SoftMinCustom : public OpDef {
public:
    explicit SoftMinCustom(const char *name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("z")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        // 定义属性
        this->Attr("dim")
            .AttrType(OPTIONAL)
            .Int(1);

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend910")
            .AddConfig("ascend310p")
            .AddConfig("ascend310b")
            .AddConfig("ascend910b");
    }
};
OP_ADD(SoftMinCustom);
} // namespace ops